# -*- coding: utf-8 -*-
# @Time    : 2021/4/29 15:56
# @Author  : huangwei
# @File    : baoxiaodan.py
# @Software: PyCharm

from bxdmethod import *
from table import get_cell
from config import args

WEIGHT_PATH = "models/table.weights"
tableNet = cv2.dnn.readNetFromDarknet(WEIGHT_PATH.replace('.weights', '.cfg'), WEIGHT_PATH)

text_sys = TextSystem(args)


def get_tabel_data(img):
    img = normal_size(img)

    img_resize_path = "img_resize.png"
    cv2.imwrite(img_resize_path, img)

    result = text_sys(img)

    # 找到表格中的所有空格
    img2 = Image.open(img_resize_path).convert('RGB')

    # row, col为表格框的最小大小
    boxes, _ = get_cell(img2, tableNet, prob=0.5, alph=10, row=50, col=20)

    # from test import draw_box
    # for box in boxes:
    #     draw_box(img, box)
    # cv2.imwrite("pic/box.png", img)
    # print("box numbers:", len(boxes))

    back_result = []  # 用于记录在表格中的文字框的序号
    # 找出每个box中包含的文字框
    sort_word = []
    for box in boxes:
        one_box = []
        for i in range(len(result)):
            dt_info = result[i]
            dt_box = dt_info['box']
            c_x = np.average(dt_box[:, 0])
            c_y = np.average(dt_box[:, 1])

            # 判断该点是否在box中
            is_in_box = in_box((c_x, c_y), box)
            if is_in_box:
                one_box.append(dt_info)
                back_result.append(i)
        if len(one_box) != 0:
            sort_word.append(one_box)

    # 判断sort_word中的项是否为上下关系
    word_list = []
    for item in sort_word:
        word = ''
        for dd in item:
            word += dd['rec'][0]

        word_list.append(word)

    # print(word_list)
    # 剩下的不在表格中的文字
    for j in range(len(result)):
        if j not in back_result:
            word = result[j]['rec'][0]
            word_list.append(word)

    # print(word_list)
    return word_list


def detect():
    img = cv2.imread("pic/code4.jpg")

    # 先确定能否识别二维码
    code_data = code_info(img)
    print(code_data)
    # table_word_list = get_tabel_data(img)

    infos = {}

    if len(code_data) != 0:
        # 存在二维码，提取出其他信息
        table_data = get_tabel_data(img)
        infos['code'] = code_data
        infos['data'] = table_data
    else:
        print("没有识别出二维码！")

    return infos


if __name__ == '__main__':
    data_dict = detect()
    print(data_dict)
